"""
Retriever（检索器）就是 "去资料库里帮你找最相关内容的工具人"。
它不回答问题，只负责「找」；找完之后，再交给大模型（比如 GPT）去「回答」。

Retriever 帮 LLM「划重点」，让 LLM 只关注最相关的信息！

工作流程：
1. 加载文档 → 2. 切分文档 → 3. 向量化存储 → 4. 创建检索器 → 5. 检索相关文档
"""

from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores import Chroma
from langchain_core.documents import Document
from langchain_huggingface import HuggingFaceEmbeddings

# 1. 加载文本 - 也可以自己改读文档
# 方法1：从文件加载（取消注释使用）
# loader = TextLoader("docs/my_notes.txt")
# docs = loader.load()

# 方法2：直接创建 Document 对象（当前使用的方式）
# Document 是 LangChain 中的基本文档单位，包含页面内容和元数据
docs = [
    Document(page_content="LangChain 是一个让大模型可以使用外部数据和工具的框架。"),
    Document(page_content="Retriever 是 LangChain 中用于从数据库中检索相关信息的组件。"),
    Document(page_content="Milvus 是一个高性能的向量数据库，用于存储和检索向量化的数据。")
]

# 2. 切分文档 - 将长文档分割成适合处理的片段
# 为什么需要切分：大模型有输入长度限制，切分后可以提高检索精度
splitter = CharacterTextSplitter(
    chunk_size=30,      # 每个文本块的最大字符数
    chunk_overlap=10    # 块之间的重叠字符数，避免信息被切断
)
split_docs = splitter.split_documents(docs)

# 打印切分结果，便于理解切分效果
print("文档切分结果：")
for i, doc in enumerate(split_docs):
    print(f"片段 {i+1}: {doc.page_content}")

# 3. 向量化 + 存储 - 将文本转换为向量并存入向量数据库
# 向量化：将文本转换为数学向量，便于计算相似度
embeddings = HuggingFaceEmbeddings(model_name="D:/models/BAAIbge-base-zh-v1.5", model_kwargs={'device': 'cpu'})


# 创建向量数据库：将切分后的文档转换为向量并存储
# Chroma 是一个轻量级的向量数据库，专门用于存储和检索向量数据
vectorstore = Chroma.from_documents(
    documents=split_docs,    # 切分后的文档
    embedding=embeddings     # 使用的嵌入模型
)

# 4. 获取 Retriever - 创建检索器实例
retriever = vectorstore.as_retriever(
    # 设置检索参数
    search_kwargs={
        "k": 1  # 返回最相关的 1 个结果
        # 其他可选参数：
        # "score_threshold": 0.8  # 相似度阈值，只返回相似度高于0.8的结果
    }
)

# 5. 检索测试 - retriever 只负责"找"，不负责"回答"
print("\n开始检索...")
query = "什么是 LangChain？"
results = retriever.get_relevant_documents(query)

print(f"查询问题: {query}")
print("检索结果:")
for i, r in enumerate(results):
    print(f"结果 {i+1}: {r.page_content}")
    print(f"元数据: {r.metadata}")  # 显示文档的元数据信息

# 后续步骤：将检索结果传递给 LLM 生成最终答案
"""
检索完成后，通常的下一步是：
1. 将检索到的文档内容作为上下文
2. 结合用户问题构造 prompt
3. 发送给 LLM 生成最终答案

示例：
context = "\n".join([r.page_content for r in results])
prompt = f"基于以下信息回答问题：\n{context}\n\n问题：{query}"
answer = llm.invoke(prompt)  # 调用大模型生成答案
"""